'''给新数据时候测试的'''
import os
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import numpy as np 
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from DataSet  import MyDataSet, ToTensor, Resize
from net import ResNet18
import torchvision.models as models
import shutil

if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    test_dataset = MyDataSet(root_dir = '/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/crop-classify/crop-all',\
         transform = transforms.Compose([
            transforms.Resize((224,224)),
            transforms.ToTensor(),
            transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
        ]))
    testset_dataloader = DataLoader(dataset = test_dataset, batch_size=1, shuffle=False, num_workers=16, pin_memory=True)

    savePath = "/media/pc/data_1/luo/data/0723/save"

    # model = ResNet18()
    model = models.resnet50(pretrained=False)
    model.fc=nn.Linear(2048,5)
    model.load_state_dict(torch.load('/home/xys/CloundShiProjects/traffic_light/trafficlight_classify/model/19lei_96_选定最终结果/model_050.pth'))
    model.cuda()
    model.eval()
    saveFile1 = os.path.join(savePath,"dun")
    os.makedirs(saveFile1)
    saveFile2 = os.path.join(savePath,"zhan")
    os.makedirs(saveFile2)
    saveFile3 = os.path.join(savePath,"zuo")
    os.makedirs(saveFile3)
    saveFile4 = os.path.join(savePath,"shuai")
    os.makedirs(saveFile4)
    saveFile5 = os.path.join(savePath,"qita")
    os.makedirs(saveFile5)
    with torch.no_grad():
        for j, val_data in enumerate(testset_dataloader):
            val_inputs, val_labels, image_path = val_data
            val_inputs = Variable(val_inputs.cuda())
            val_labels = Variable(val_labels.cuda())
            # val_inputs, val_labels = Variable(val_data["image"].cuda()).float(), Variable(val_data["label"].cuda())
            val_outputs = model(val_inputs)
            val_outputs.detach_()
            _, val_predicted = torch.max(val_outputs.data, 1)
            # print(val_predicted)
            # print(image_path)
            if(val_predicted[0] == 0):  
                shutil.copy2(image_path[0], saveFile1)
            elif(val_predicted[0] == 1):
                shutil.copy2(image_path[0], saveFile2)
            elif(val_predicted[0] == 2):
                shutil.copy2(image_path[0], saveFile3)
            elif(val_predicted[0] == 3):
                shutil.copy2(image_path[0], saveFile4)
            else:
                shutil.copy2(image_path[0], saveFile5)
            